#pragma once

#include <vector>
#include <algorithm>
#include <armadillo>

#include <boost/archive/binary_oarchive.hpp>
#include <boost/archive/binary_iarchive.hpp>
#include <boost/serialization/split_free.hpp>

#include <boost/mpi.hpp>

#include "helper.h"

using namespace std;
using namespace arma;
namespace mpi = boost::mpi;

#define MPI_IS_MASTER (this->world.rank() == 0)
#define MPI_SIZE (SIZE(this->world))
#define MPI_RANK (this->world.rank())
#define MPI_WORLD (this->world)

namespace boost { 
  namespace serialization {
    template<class Archive>
    void save(Archive & ar, const mat & A, unsigned int version)
    {
      
      int m = A.n_rows, n = A.n_cols;
      const double* A_mem = A.memptr();
      std::vector<double> A_vec(A_mem, A_mem+A.n_elem);
      //cout << "SAVE m=" << m << " n=" << n << " sz=" << SIZE(A_vec) << endl;
      ar & m;
      ar & n;
      ar & A_vec;
    }

    template<class Archive>
    void load(Archive & ar, mat & A, unsigned int version)
    {
      int m, n;
      std::vector<double> A_vec;
      ar & m;
      ar & n;
      ar & A_vec;
      //cout << "LOAD m=" << m << " n=" << n << " sz=" << SIZE(A_vec) << endl;
      A = mat(A_vec.data(), m, n);
    }
  }
}

BOOST_SERIALIZATION_SPLIT_FREE(mat)

#define PART_BROADCAST_TAG 0
#define PART_GATHER_TAG 1

template<typename T>
class PART {
 public:
  static void cat(const vector<T>&, T&);
  static void split(const T&, const vector<int>&, vector<T>&);
  static void split_col(const T&, const vector<int>&, vector<T>&);

  static void broadcast(mpi::communicator& world, const vector<T>& v, T& holder) {
    int n = min(SIZE(world), SIZE(v));
    if (world.rank() == 0) {
      vector<mpi::request> reqs;
      FORE(i, 1, n) reqs.push_back(world.isend(i, PART_BROADCAST_TAG, v[i]));
      holder = v[0];
      wait_all(ALL(reqs));
    } else if (world.rank() < n) {
      mpi::request req = world.irecv(0, PART_BROADCAST_TAG, holder);
      req.wait();
    }
  }
};

template<>
class PART<mat> {
 public:
  static void cat(const vector<mat>& v, mat& a) {
    if (SIZE(v) == 0) {
      a = mat();
      return;
    }
    
    int s = 0;
    FORE(i, 0, SIZE(v)) s += v[i].n_rows;
    a = mat(s, v[0].n_cols);
    int r = 0;
    FORE(i, 0, SIZE(v)) {
      a.rows(r, r+v[i].n_rows-1) = v[i];
      r += v[i].n_rows;
    }
  }

  static void split_idx(int n, int n_part, vector<int>& r) {
    int part_size = n / n_part + 1;
    int f = 0;
    r = vector<int>(n_part);
    FORE(i, 0, n_part) {
      r[i] = f;
      f += (i < n % n_part ? part_size : part_size-1);
    }
  }

  static void split(const mat& A, const vector<int>& r, vector<mat>& v) {
    v = vector<mat>(SIZE(r));
    int n = SIZE(r);
    FORE(i, 0, n) {
      v[i] = A.rows(r[i], (i < n-1) ? r[i+1]-1 : A.n_rows-1);
    }
  }

  static void split_col(const mat& A, const vector<int>& r, vector<mat>& v) {
    v = vector<mat>(SIZE(r));
    int n = SIZE(r);
    FORE(i, 0, n) {
      v[i] = A.cols(r[i], (i < n-1) ? r[i+1]-1 : A.n_cols-1);
    }
  }

  static void broadcast(mpi::communicator& world, const vector<mat>& v, mat& holder) {
    if (world.rank() == 0) {
      int n = SIZE(world);
      if (n != SIZE(v)) FATAL("Check size(v) != size(world)");
      vector<mpi::request> reqs;
      FORE(i, 1, n) reqs.push_back(world.isend(i, PART_BROADCAST_TAG, v[i]));
      holder = v[0];
      wait_all(ALL(reqs));
    } else {
      mpi::request req = world.irecv(0, PART_BROADCAST_TAG, holder);
      req.wait();
    }
  }

  static void save(mpi::communicator& world, const string& prefix, const arma::mat& a, file_type type = arma_binary) {
    stringstream ss;
    ss << prefix << "_" << world.rank();
    a.save(ss.str(), type);
  }

  static void load(mpi::communicator& world, const string& prefix, arma::mat& a, file_type type = arma_binary) {
    stringstream ss;
    ss << prefix << "_" << world.rank();
    a.load(ss.str(), type);
  }
};
